# -*- coding:utf-8 -*-
import re
import numpy as np
from metlib.datetime import T, TD
from metlib.kits import *
from weblib.common.file_transporter import *
from filecache.filecache_manager import *
from matplotlib.cm import *

# merra_transporter = FileTransporter(remote_location='s3://windres-datasets', retry=2, retry_interval=1)
# merra_fc_manager = FileCacheManager(bucket='merra_raw',
#                                     transporter=merra_transporter)
# merra2d_fc_manager = FileCacheManager(bucket='merra2d_raw',
#                                     transporter=merra_transporter)
# merra_kmz_manager = FileCacheManager('merra_kmz_img_tmp')

merra_mean_varnames = ['ws50m', 'wpd', 't10m', 'slp', 'ps' ]
merra_pt_varnames = ['ws50m', 'wd50m', 'wpd', 't10m',  'slp', 'ps' ]
merra_py_varnames = ['mean', 'rose', 'wpdrose', 'dist']
merra_ry_varnames = ['ws50m', 'wpd', 't10m', 'slp', 'ps', ]
merra_pm_varnames = merra_mean_varnames
merra_pd_varnames = merra_mean_varnames

merra_years = [str(y) for y in range(1979, 2015)]

merra_full_lons = np.arange(-180.0, 179.34, 2.0/3)
merra_full_lats = np.arange(-90.0, 90.1, 0.5)

merra_zh_dataset_d = {
    'merra': 'Merra'
}

merra_zh_subset_d = {
   "PT": u"时间序列",
   "PY": u"年统计值",
   "PM": u"月统计值",
   "PD": u"日变化",
   "RY": u"年统计值",
}

merra_zh_varname_d = {
    "ws50m": u"风速(50m)",
    "wd50m": u"风向(50m)",
    "wpd": u"风功率密度",
    "t10m": u"温度(10m)",
    "t2m": u"温度(2m)",
    "ts": u"温度(地面)",
    "slp": u"海平面气压",
    "ps": u"地面气压",
    "disph": u"位移高度",
    "mean": u"汇总统计",
    "rose": u"风玫瑰",
    "windrose": u"风玫瑰",
    "wpdrose": u"风功率玫瑰",
    "dist": u"风速分布",
}

# 单位的字典
merra_units_d = {
    "ws50m": u"m/s",
    "wd50m": u"°",
    "wpd": u"W/m²",
    "t10m": u"°C",
    "t2m": u"°C",
    "ts": u"°C",
    "ps": u"hPa",
    "slp": u"hPa",
    "disph": u"m",
    "summary": u"",
    "rose": u"‰",
    "windrose": u"‰",
    "wpdrose": u"kWh/m²",
    "dist": u"%",
}

merra_suggest_range_d = {
    "ws50m": (3.0, 10.0),
    "wd50m": (0.0, 360.0),
    "wpd": (0.0, 1000.0),
    "t10m": (-10.0, 30.0),
    "t2m": (-10.0, 30.0),
    "ts": (-10.0, 30.0),
    "ps": (None, None),
    "slp": (1000, 1020),
}

merra_cmap_d = {
    'jet': 'jet',
    'YlGnBu': 'YlGnBu',
    'rainbow': 'gist_rainbow_r',
    'Spectral_r': 'Spectral_r',
    'Spectral': 'Spectral',
}

# 推荐的色表
merra_suggest_cmap_d = {
    't10m': 'temperature',
    't2m': 'temperature',
    'ts': 'temperature',
    'ps': 'pressure',
    'slp': 'pressure',
}

def latlon_to_merra_latlon(lat, lon):
    mlat = np.round(lat * 2) / 2
    mlon = np.round(lon * 3.0/2.0) / 3.0 * 2.0
    return '%.1f_%.1f' % (mlat, mlon)

def get_merra_rect(lon1, lat1, lon2, lat2):
    latmin = min(lat1, lat2); latmax = max(lat1, lat2)
    lonmin = min(lon1, lon2); lonmax = max(lon1, lon2)
    latmin -= 0.5; latmax += 0.5
    lonmin -= 2.0 / 3; lonmax += 2.0 / 3
    mlatmin, mlonmin = latlon_to_merra_latlon(latmin, lonmin).split('_')
    mlatmax, mlonmax = latlon_to_merra_latlon(latmax, lonmax).split('_')

    res = {
        'lon1': mlonmin,
        'lat1': mlatmin,
        'lon2': mlonmax,
        'lat2': mlatmax,
        'jy_ix': '%s:%s_%s:%s' % (mlatmin, mlatmax, mlonmin, mlonmax)
    }
    return res

def parse_merra_uri(uri):
    result = {}
    m = re.match(r'(?P<dataset>[^/]+)/(?P<subset>[^/]+)/(?P<varname>[^/]+)/(?P<time>[^-+Z/]+)(Z?)(?P<timezone>[-+0-9]*)/(?P<lat_lon>[^/]+)$', uri)
    if not m:
        raise ValueError('%s is not a proper merra uri' % uri)
    for field in ['dataset', 'subset', 'varname', 'time', 'lat_lon']:
        result[field] = m.group(field)

    time_str = m.group('time')
    timezone_str = m.group('timezone')
    try:
        timezone = int(timezone_str)
    except ValueError:
        timezone = 0

    if ',' in time_str:
        times = time_str.split(',')
        result['times'] = times
        result['timetype'] = 'year_list'
    elif ':' in time_str:
        begdt, enddt, tdelta = parse_slice_str(time_str, default_step='1h')
        begdt = T(begdt)
        enddt = T(enddt)
        result['begdt'] = begdt
        result['enddt'] = enddt
        result['tdelta'] = tdelta
        result['timetype'] = 'dtrange'
    elif len(m.group('time')) == 4:  # year
        begdt = T(m.group('time')+'0101')
        enddt = begdt + TD('1Y')
        result['begdt'] = begdt
        result['enddt'] = enddt
        result['tdelta'] = '1h'
        result['timetype'] = 'year'
    else:
        raise ValueError('%s is not a proper merra uri' % uri)
    result['timezone'] = timezone
    result['subset_long'] = {'PT': 'PointTimeseries'}.get(result['subset'], result['subset'])
    lat, lon = m.group('lat_lon').split('_')
    if ':' in lat:
        lat1, lat2 = lat.split(':')
        lon1, lon2 = lon.split(':')
        lat1 = float(lat1); lat2 = float(lat2)
        lon1 = float(lon1); lon2 = float(lon2)
        mlat1, mlon1 = latlon_to_merra_latlon(lat1, lon1).split('_')
        mlat2, mlon2 = latlon_to_merra_latlon(lat2, lon2).split('_')
        jy_ix = '%s:%s_%s:%s' % (mlat1, mlat2, mlon1, mlon2)
    else:
        lat, lon = [float(v) for v in m.group('lat_lon').split('_')]
        jy_ix = latlon_to_merra_latlon(lat, lon)
    result['jy_ix'] = jy_ix

    return result

def merra_lookup_lonlat(lon, lat):
    jy_ix = latlon_to_merra_latlon(lat, lon)
    mlat = np.round(lat * 2) / 2
    mlon = np.round(lon * 3.0/2.0) / 3.0 * 2.0
    point_info = {
        "grid_lon": mlon,
        "grid_lat": mlat,
        "jy_ix": jy_ix,
    }
    return point_info

def merra_yearlymean2summary(orig):
    tags = orig['tags']
    sub_varnames = ['ws50m', 'wpd', 't10m', 't2m', 'ts', 'ps', 'slp']
    info = {
        'sub_varnames': sub_varnames,
        'type': 'merra_summary',
    }
    data = {}
    for vn in sub_varnames:
        rawvar = orig['contents'][vn]
        subvar = {
            'name': vn,
            'zh_name': merra_zh_varname_d.get(vn, vn),
            'values': rawvar['data']['values'],
            'std': rawvar['data']['std'],
            'count': rawvar['data']['count'],
            'units': rawvar['data']['units'],
        }
        data[vn] = subvar
    common = orig['common'].copy()
    common['coords'] = {}
    for key in ['dataset', 'subset', 'time', 'jy_ix']:
        common['coords'][key] = common['info'].pop(key, '')

    res = {
        'uri': orig['uri'],
        'common': common,
        'tags': {
            'data': orig['uri']
        },
        'contents': {
            'data': {
                'info': info,
                'data': data,
                'coords': {
                    'varname': 'mean',
                }
            },
        }
    }

    return res
